from text_extraction_tool import TextExtractionTool
from vector_db import MyVectorDBConnector
from embeddings import get_embeddings
from openai_client import get_completion
from english_utils import split_en_text
from chinese_utils import split_zh_text


class RAG_Bot:
    def __init__(self, vector_db, llm_api, n_results=2):
        self.vector_db = vector_db
        self.llm_api = llm_api
        self.n_results = n_results

    def chat(self, user_query):
        # 1. 检索
        search_results = self.vector_db.search(user_query, self.n_results)

        # 2. 构建 Prompt
        prompt = build_prompt(
            prompt_template, info=search_results['documents'][0], query=user_query)

        # 3. 调用 LLM
        response = self.llm_api(prompt)
        return response


prompt_template = """
你是一个问答机器人。
你的任务是根据下述给定的已知信息回答用户问题。
确保你的回复完全依据下述已知信息。不要编造答案。
如果下述已知信息不足以回答用户的问题，请直接回复"我无法回答您的问题"。

已知信息:
__INFO__

用户问：
__QUERY__

请用中文回答用户问题。
"""


# 构建 Prompt 模板
def build_prompt(prompt_template, **kwargs):
    """将 Prompt 模板赋值"""
    prompt = prompt_template
    for k, v in kwargs.items():
        if isinstance(v, str):
            val = v
        elif isinstance(v, list) and all(isinstance(elem, str) for elem in v):
            val = '\n'.join(v)
        else:
            val = str(v)
        prompt = prompt.replace(f"__{k.upper()}__", val)
    return prompt


# 创建一个向量数据库对象
vector_db = MyVectorDBConnector("chromadb", get_embeddings)

# PDF 文档加载与切割（为了演示方便，我们只取两页）
# paragraphs = TextExtractionTool.extract_text_from_pdf("../document/学习笔记.pdf", page_numbers=[10, 11, 12], min_line_length=10)
paragraphs = """
## MyBatis XML中解决尖括号导致报错的两种常见方法
MyBatis是一个持久层框架，通过XML文件配置SQL语句与Java方法的映射关系。在XML中编写SQL时，有时候会遇到尖括号（<）导致报错的问题。为了解决这个问题，有两种常见的方法。
首先，可以使用 `<![CDATA[ … ]]>` 将SQL语句包裹起来，这样可以告诉解析器将其中的内容视为字符数据而不是XML标签。
其次，可以使用替换符号的方式，即将 `<` 替换为 `&lt;`，将 `>` 替换为 `&gt;`。这样可以避免XML解析错误，确保SQL语句在XML中正确无误地被解析执行。选择哪种方法取决于个人偏好和项目的需求。
### 使用 `<![CDATA[ … ]]>` 包裹SQL语句
在XML中，`<![CDATA[ … ]]>` 用于将其中的内容视为字符数据而不是XML标签。因此，可以将SQL语句包裹在 `<![CDATA[ … ]]>` 中，这样就可以避免尖括号导致的XML解析错误。示例如下：

```xml
<select id="selectUser" resultType="User">
    select * from user where id <![CDATA[<]]> #{id}
</select>
```

### 使用替换符号
在XML中，使用 `&lt;` 代替 `<`，使用 `&gt;` 代替 `>`，这样可以避免XML解析错误。示例如下：
```xml
<select id="selectUser" resultType="User">
    select * from user where id &lt; #{id}
</select>
```

### MySQL字符串操作
1. **使用replace()函数替换指定字段中的字符串**
   在MySQL中，可以使用`replace()`函数来替换指定字段中的字符串。例如，下面的语句将表中某个字段中的'oldVal'替换为'newVal':

   ```sql
   update 表 set 字段 = replace(字段, 'oldVal', 'newVal');
   ```

2. **使用locate()函数判断字段是否包含某个字符**
   通过`locate()`函数可以判断字段是否包含某个字符。下面的示例展示了如何使用`locate()`函数来判断字段是否包含特定字符：

   ```sql
   -- 不包含
   select * from 表 where locate('某字符', 字段) = 0;

   -- 包含
   select * from 表 where locate('某字符', 字段) > 0;
   ```

### 更新SQL查询模式
1. **查询当前SQL模式**
   若要查询当前的SQL模式，可以使用以下语句：
   ```sql
   SELECT @@sql_mode;
   ```

2. **更新SQL查询模式**
   要更新SQL模式，可以使用`SET @@sql_mode`语句，例如：
   ```sql
   SET @@sql_mode = 'ONLY_FULL_GROUP_BY,STRICT_TRANS_TABLES,NO_ZERO_IN_DATE,NO_ZERO_DATE,ERROR_FOR_DIVISION_BY_ZERO,NO_ENGINE_SUBSTITUTION';
   ```

3. **使用更新后的SQL模式进行查询**
   此外，下面的查询示例展示了如何使用更新后的SQL模式进行查询：
   ```sql
   SELECT * FROM tjw_agency_account aa GROUP BY aa.`agt_id`;
   ```

"""
print("PDF 文档内容", paragraphs)
# chunks = to_keywords(paragraphs[0])  # 文本分割的粒度处理
chunks = split_zh_text(paragraphs)  # 文本分割的粒度处理
print("分割后的文本", chunks)
# 向向量数据库中添加文档
vector_db.add_documents(chunks)

# paragraphs = TextExtractionTool.extract_text_from_txt("../document")
# # 向向量数据库中添加文档
# vector_db.add_documents(paragraphs[0]['content'])


# 创建一个RAG机器人
bot = RAG_Bot(
    vector_db,
    llm_api=get_completion
)

# user_query = "llama 2的参数？"

# user_query = "用好 AI 的核心心法是什么？"
# user_query = "⼤模型可以用来做什么？它的应⽤场景有哪些" # 回答不好
# user_query = "什么是 AI 应用？"
# user_query = "基于机器学习、神经⽹络的是什么类型的应用？"
# user_query = "优秀作品展示"

user_query = "提示⼯程的目的是什么？"
user_query = "mybatis 大于号处理，并举例说明"
response = bot.chat(user_query)
print("===回复===")
print(response)
